#Author: Tim Henderson
#Maintainer: Tim Henderson @ 
#   tim.tadh@gmail.com
#   400 Braemer Ct. Gahanna OH 43230 USA
#Purpose: To allow the user to save files in an encrypted archive
#Copyright: (c)2006 Tim Henderson. All Rights Reserved. This module is part of 
#   the program secureData.
#License: secureData and all of its component parts are licensed under the 
#   terms and conditions of the GNU General Public License Version 2. Later versions
#   of this license do not apply to secureData or its component parts. If you have 
#   not recieved a copy of the license with your software please contact the Free 
#   Software Foundation at:
#       51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA
#   or you may also locate the license at the following internet address
#       http://www.gnu.org/licenses/gpl.txt

from Crypto.Cipher import AES
from Crypto.Hash import MD2, MD5, RIPEMD, SHA256
from Crypto.Util import randpool
#import numpy.random.mtrand as mtrand
import random
import os, stat, time
import nDDB
import cStringIO
import glob

class EFAerror(Exception): pass
class ControlIndexTooLongError(EFAerror): pass
class FileAlreadyInArchiveError(EFAerror): pass
class FileDoesntExistError(EFAerror): pass
class FileNotInArchiveError(EFAerror): pass
class BadControlDDB(EFAerror): pass
class InvalidArchiveError(EFAerror): pass

class RandDataGenError(Exception): pass

KEY_GEN_BLOCK = 1024

def hMD2gen(s, prounds=0):
    """use of this functions is deprecated due to security concerns"""
    md2 = MD2.new()
    md2.update(s)
    for x in range(prounds+1): md2.update(md2.digest())
    while True:
        yield md2.digest()
        md2.update(md2.digest())

def MD2gen(s, prounds=0):
    """use of this functions is deprecated due to security concerns"""
    md2 = MD2.new()
    md2.update(s)
    for x in range(prounds+1): md2.update(md2.digest())
    while True:
        yield md2.digest()
        md2.update(md2.digest())

def hSHA256gen(s, prounds=0):
    sha = SHA256.new()
    sha.update(s)
    for x in range(prounds+1): sha.update(sha.digest())
    while True:
        yield sha.hexdigest()
        sha.update(sha.digest())

def SHA256gen(s, prounds=0):
    sha = SHA256.new()
    sha.update(s)
    for x in range(prounds+1): sha.update(sha.digest())
    while True:
        yield sha.digest()
        sha.update(sha.digest())

def RMDhexGen(s, prounds=0):
    """use of this functions is deprecated due to security concerns"""
    rmd = RIPEMD.new()
    rmd.update(s)
    for x in range(prounds+1): rmd.update(rmd.digest())
    while True:
        yield rmd.hexdigest()
        rmd.update(rmd.digest())

def __updateHash(f, rlen, mtprng, urandom, word_size, hfun1, hfun2):
    '''helper function for SHA256gen_file_prng'''
    r = ''.join([str(chr(mtprng(0, 256))) for x in range(word_size)])
    rsha = SHA256.new()
    rsha.update(r) #ensures prng's state is not revealed
    r = rsha.digest()
    del rsha
    
    u = urandom(32)
    
    s = f.read(word_size)
    
    hfun2.update(s)
    
    hfun1.update(s)
    hfun1.update(r)
    hfun1.update(u)
    hfun1.update(hfun1.digest())
    hfun1.update(hfun2.digest())

def SHA256gen_file_prng(filename, prounds=0, **kwargs):
    global KEY_GEN_BLOCK
    block = KEY_GEN_BLOCK
    
    f = open(filename, 'rb')
    flen = int(os.stat(filename)[stat.ST_SIZE])
    
    #standard python Mersenne twister prng
    if flen - block > 0: a = block
    else:  a = flen
    random.seed(SHA256gen(f.read(a), 32).next())
    f.seek(0)
    random.getrandbits(block * 4)
    
    sha = SHA256.new()
    rmd = RIPEMD.new()
    
    read = 0
    while flen - read > block:
        __updateHash(f, block, random.randrange, os.urandom, 32, sha, rmd)
        update_progress(**kwargs)
        read += block
    
    __updateHash(f, (flen - read), random.randrange, os.urandom, 64, sha, rmd)
    update_progress(**kwargs)
    
    for x in range(prounds+1): sha.update(sha.digest())
    while True:
        yield sha.digest()
        sha.update(sha.digest())


def makeFileQueue(filename):
    """use of this functions is deprecated due to security concerns"""
    f = open(filename, 'rb')
    global KEY_GEN_BLOCK
    block = KEY_GEN_BLOCK
    tread = 128**2/2
    flen = int(os.stat(filename)[stat.ST_SIZE])
    if flen < tread: tread = flen
    read = 0
    queue = []
    while tread - read > block:
        queue.append(f.read(block))
        read += block
    queue.append(f.read(tread - read))
    f.close()
    return queue

def randData_RIPEMDgen(queue): #returns 160 bit hexadecimal digest
                               #The stack must contain indexed container objects
    """use of this functions is deprecated due to security concerns"""
    if len(queue) < 1:
        raise RandDataGenError, "queue is empty"
    elif len(queue) < 2:
        rmd = RIPEMD.new()
        s = queue[0]
        del queue[0]
        rmd.update(s)
        rmd.update(rmd.digest())
        yield rmd.digest()
        raise StopIteration, "stack of strings is too small for this generator"
    
    rmd = RIPEMD.new()
    s = queue[0]
    del queue[0]
    rmd.update(s)
    rmd.update(rmd.digest())
    c = rmd.digest()
    
    while True:
        yield c
        if len(queue) < 1: 
            raise StopIteration, "the stack is empty"
        s = queue[0]
        del queue[0]
        rmd.update(s)
        rmd.update(rmd.digest())
        c = rmd.digest()
        
def update_progress(wieght=1, **kwargs):
    if kwargs.has_key('progress'):
        kwargs['progress'].increment(wieght)
 
class EncryptedFA(object):
    
    CLM = 12
    BLOCK = 2**20
    
    def __init__(self, path):
        self.path = os.path.abspath(path)
        if self._pathExists(path):
            self._readIn()
        else:
            self._new()
    
    def _pathExists(self, path):
        if len(glob.glob(path)) > 0:
            return True;
        else:
            return False;
    
    def _readIn(self):
        fmrb = open(self.path, 'rb')
        try:
            self.controlLen = self._readControlLen(fmrb)
        except:
            fmrb.close()
            self.controlLen = None
            self.control = None
            raise InvalidArchiveError, "The CLM mark was unreadable"
        self.control = self._readControlDDB(self.controlLen, fmrb)
        fmrb.close()
    
    def _readControlLen(self, f):
        mylen = int(os.stat(self.path)[stat.ST_SIZE])
        f.seek(mylen-self.CLM)
        slen = f.read(self.CLM)
        return int(slen)
    
    def _readControlDDB(self, clen, f):
        mylen = int(os.stat(self.path)[stat.ST_SIZE])
        f.seek(mylen - self.CLM - clen)
        ddb = f.read(clen)
        try:
            d = nDDB.decodeDDB(ddb)
            if d == None:
                raise BadControlDDB, 'The control index was corrupted'
            return d
        except:
            raise BadControlDDB, 'The control index was corrupted'
            return None
    
    def getFileList(self):
        return self.control.keys()
    
    def getFileInfo(self, filename):
        if self.fileInArchive(filename):
            return self.control[filename]
    
    def fileInArchive(self, filename):
        return self.control.has_key(filename)
        
    def readFileEncrypted_interatively(self, filename):
        #this is meant to be used with addIterFile_preEncrypted
        iF = open(self.path, 'rb')
        if self.control.has_key(filename):
            flen = int(self.control[filename]['len'])
            if flen < self.BLOCK:
                
                yield self._readSpot(int(self.control[filename]['start']),
                                        flen, iF)
            else:
                clen = flen
                tread = 0
                while clen - self.BLOCK > 0:
                    yield self._readSpot(int(self.control[filename]['start']) + tread,
                                        self.BLOCK, iF)
                    tread += self.BLOCK
                    clen -= self.BLOCK
                
                yield self._readSpot(int(self.control[filename]['start']) + tread,
                                        clen, iF)
        else:
            raise FileNotInArchiveError, "The file %s is not in the archive" % filename
    
    def readFile(self, oF, name, keyGen, added, **kwargs):
        added = int(added)
        iF = open(self.path, 'rb')
        if self.control.has_key(name):
            flen = int(self.control[name]['len'])
            if flen < self.BLOCK:
                
                aes = AES.new(keyGen.next(), AES.MODE_ECB)
                spot = aes.decrypt(self._readSpot(int(self.control[name]['start']),
                                        flen, iF))
                if added != 0:
                    oF.write(spot[:-added])
                else:
                    oF.write(spot)
                del aes
            else:
                clen = flen
                tread = 0
                while clen - self.BLOCK > 0:
                    aes = AES.new(keyGen.next(), AES.MODE_ECB)
                    oF.write(aes.decrypt(self._readSpot(int(self.control[name]['start']) + tread,
                                        self.BLOCK, iF)))
                    del aes
                    tread += self.BLOCK
                    clen -= self.BLOCK
                    update_progress(**kwargs)
                aes = AES.new(keyGen.next(), AES.MODE_ECB)
                spot = aes.decrypt(self._readSpot(int(self.control[name]['start']) + tread,
                                        clen, iF))
                if added != 0:
                    oF.write(spot[:-added])
                else:
                    oF.write(spot)
                del aes
            update_progress(**kwargs)
        else:
            raise FileNotInArchiveError, "The file %s is not in the archive" % name
    
    def _readSpot(self, start, slen, f):
        f.seek(start)
        spot = f.read(slen)
        return spot
    
    def _resetControlLen(self):
        self.controlLen = len(nDDB.makeAdvanceDDB(self.control))
    
    def _new(self):
        fmwb = open(self.path, 'wb')
        self.control = {}
        self.controlLen = len(nDDB.makeAdvanceDDB(self.control))
        fmwb.close()
        self._writeControl()
    
    def __len__(self):
        return self.__customControlLen(self.controlLen)
    
    def __customControlLen(self, clen):
        length = self.CLM
        length += self.controlLen
        for name in self.control:
            length += int(self.control[name]['len'])
        return length
    
    def _shiftLeft(self, fmrbp, start, amt, tlen, **kwargs):
        wcur = start
        rcur = wcur + amt
        fmrbp.seek(rcur)
        ch = fmrbp.read(self.BLOCK)
        while len(ch) == self.BLOCK:
            fmrbp.seek(wcur)
            fmrbp.write(ch)
            fmrbp.flush()
            wcur += self.BLOCK
            rcur += self.BLOCK
            fmrbp.seek(rcur)
            ch = fmrbp.read(self.BLOCK)
            update_progress(**kwargs)
            
        fmrbp.seek(wcur)
        fmrbp.write(ch)
        fmrbp.flush()
        update_progress(**kwargs)
        
        if tlen != None:
            #print 'TRUNCATED', tlen
            fmrbp.truncate(tlen)
    
    def _shiftRight(self, fmrbp, start, amt, tlen):
        if tlen - start < self.BLOCK: block = tlen - start
        else: block = self.BLOCK
        rcur = tlen
        wcur = rcur + amt
        ch = ''
        while rcur - block > start:
            rcur -= block
            wcur -= block
            
            fmrbp.seek(rcur)
            ch = fmrbp.read(block)
            fmrbp.seek(wcur)
            fmrbp.write(ch)
            fmrbp.flush()
            
        a = rcur - start
        rcur = start
        wcur = start + amt
        
        fmrbp.seek(rcur)
        ch = fmrbp.read(a)
        fmrbp.seek(wcur)
        fmrbp.write(ch)
        fmrbp.flush()
    
    def _deleteFileFromControl(self, fname, **kwargs):
        flen = int(self.control[fname]['len'])
        fstart = int(self.control[fname]['start'])
        
        olen = len(self)
        olenslen = len(str(olen))
        oclen = self.controlLen
        mlen = len(self.control[fname]['start'])
        
        del self.control[fname]
        
        for key in self.control:
            cs = int(self.control[key]['start'])
            if cs > fstart:
                self.control[key]['start'] = str(cs - flen)
        
        self._resetControlLen()
    
    def deleteFile(self, fname, **kwargs):
        if self.control.has_key(fname):
            flen = int(self.control[fname]['len'])
            fstart = int(self.control[fname]['start'])
            
            oclen = self.controlLen
            olen = len(self)
            self._deleteFileFromControl(fname)
            update_progress(**kwargs)
            fmrbp = open(self.path, 'rb+')
            #print fstart, flen, len(self)
            
            self._shiftLeft(fmrbp, fstart, flen, (olen - oclen - self.CLM - flen), **kwargs)
            self._writeControl(fmrbp)
            
            fmrbp.close()
            update_progress(**kwargs)
            print "DELETE FROM EFA"#, len(self), int(os.stat(self.path)[stat.ST_SIZE])
        else:
            raise FileNotInArchiveError, "The file %s is not in the archive" % fname
    
    def _addFileToControl(self, fname, flen):
        # flen = int(os.stat(abspath)[stat.ST_SIZE])
        # if flen%16 !=0:
            # flen += (16 - flen%16)
        # fname = os.path.basename(abspath)
        
        olen = len(self)
        olenslen = len(str(olen)) + 2
        oclen = self.controlLen + self.CLM
        
        if self.control.has_key(fname): 
            raise FileAlreadyInArchiveError, 'You already have a file of the name %s in the EFA' % (fname,)
        
        self.control.update({fname:{'len':flen, 'start':str(olen - oclen)}})
        
        self._resetControlLen()
    
    def addIterFile_preEncrypted(self, iFi, fname, flen, **kwargs):
        #this is meant to be used with readFileEncrypted_interatively
        if flen%16 !=0:
            flen += (16 - flen%16)
        oldLen = self.CLM + self.controlLen
        tOldLen = len(self)
        self._addFileToControl(fname, flen)
        
        fmrbp = open(self.path, 'rb+')
        
        fmrbp.seek(tOldLen - oldLen)
        
        for spot in iFi:
            fmrbp.write(spot)
            fmrbp.flush()
        
        self._writeControl(fmrbp)
        fmrbp.close()
    
    def addFile(self, iF, fname, flen, keyGen, **kwargs): #returns an integer added == spaces added
        #this needs to be changed to reading in by blocks bwe
        # path = os.path.abspath(path)
        # fname = os.path.basename(path)
        # 
        # if not self._pathExists(path):
            # raise FileDoesntExistError, 'The path %s does not exist' % (path,)
        
        #flen = int(os.stat(path)[stat.ST_SIZE])
        if flen%16 !=0:
            flen += (16 - flen%16)
        oldLen = self.CLM + self.controlLen
        tOldLen = len(self)
        self._addFileToControl(fname, flen)
        
        fmrbp = open(self.path, 'rb+')
        #print flen
        
        fmrbp.seek(tOldLen - oldLen)        
        
        #iF = open(path, 'rb')
        if flen < self.BLOCK:
            aes = AES.new(keyGen.next(), AES.MODE_ECB)
            spot, added = self._appendSpaces(iF.read())
            fmrbp.write(aes.encrypt(spot))
            del aes
            fmrbp.flush()
        else:
            clen = flen
            while clen - self.BLOCK > 0:
                aes = AES.new(keyGen.next(), AES.MODE_ECB)
                fmrbp.write(aes.encrypt(iF.read(self.BLOCK)))
                del aes
                fmrbp.flush()
                clen -= self.BLOCK
                update_progress(1024, **kwargs)
            aes = AES.new(keyGen.next(), AES.MODE_ECB)
            spot, added = self._appendSpaces(iF.read(clen))
            fmrbp.write(aes.encrypt(spot))
            del aes
            fmrbp.flush()
            
        update_progress(1025, **kwargs)
        
        self._writeControl(fmrbp)
        fmrbp.close()
        
        update_progress(**kwargs)
        
        return added
    
    def _appendSpaces(self, plaintext):
        x = 0
        if len(plaintext)%16 != 0:
            x = 16 - len(plaintext)%16
            plaintext += ' '*x
        return plaintext, x
    
    def _writeControl(self, fmrbp=None):
        if fmrbp == None: fmrbp = open(self.path, 'rb+')
        cddb = nDDB.makeAdvanceDDB(self.control)
        cddbLen = str(len(cddb))
        
        if len(cddbLen) > self.CLM: 
            raise ControlIndexTooLongError, 'The control index has exceeded the allotted space of %s bytes' % ((self.CLM * '9'),)
        
        cddbLen = (self.CLM - len(cddbLen))*'0' + cddbLen
        
        mylen = len(self) - self.CLM - self.controlLen
        fmrbp.seek(mylen)
        fmrbp.write(cddb)
        fmrbp.flush()
        fmrbp.seek(mylen + len(cddb))
        fmrbp.write(cddbLen)
        fmrbp.flush()
        
        if fmrbp == None: fmrbp.close()